import os
import sys
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../'))
import time
import cv2
import numpy as np
import torch
import torch.utils.data
from vortex.models.yolov3 import YoloBody
from vortex.network.head.yolov3 import YoloLoss
from vortex.data.voc import PascalVOCDataset
from vortex.utils.misc import AverageMeter
from vortex.utils.engine import save_state_dict
from vortex.engine.yolov3 import get_transforms


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class YoloV3Trainer(object):
    def __init__(self):
        # load model 
        anchors = np.array([[[116, 90], [156, 198], [373, 326]],
                    [[30, 61], [62, 45], [59, 119]],
                    [[10, 13], [16, 30], [33, 23]]])
        num_classes = 20
        image_size = (416, 416)
        self.model = YoloBody(anchors, num_classes)
        self.model = self.model.to(device)
        pretrained_path = './checkpoints/yolov3_pretrain.pth'
        state_dict = torch.load(pretrained_path, map_location='cpu')
        self.model.load_state_dict(state_dict, strict=False)

        # load loss
        self.yolo_losses = []
        for i in range(3):
            self.yolo_losses.append(YoloLoss(anchors[i], num_classes, image_size))

        # optimizer
        lr = 0.001
        momentum = 0.9
        weight_decay = 0.0005
        self.optimizer = torch.optim.SGD(
            params=self.model.parameters(), 
            lr=lr, momentum=momentum, weight_decay=weight_decay
        )
        
        # load dataset
        data_root = r'F:\Data\VOCdevkit\VOC2007'
        transforms = get_transforms('train')
        train_dataset = PascalVOCDataset(data_root, 'train', transforms, bg0=False)
        batch_size = 2
        nworkers = 4
        self.train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, 
                                                        shuffle=True, collate_fn=train_dataset.collate_fn, 
                                                        num_workers=nworkers, pin_memory=True)

        # train parameters
        self.epochs = 200
        self.print_freq = 20
        self.checkpoint_freq = 10

    def train_epoch(self, epoch):
        self.model.train()
        batch_time = AverageMeter()  # forward prop. + back prop. time
        data_time = AverageMeter()  # data loading time
        losses = AverageMeter()  # loss

        start = time.time()

        for i, (images, boxes, labels) in enumerate(self.train_loader):
            data_time.update(time.time() - start)

            images = images.to(device)
            boxes = [b.to(device) for b in boxes]
            labels = [l.to(device) for l in labels]
            
            # loss = self.model(images, boxes, labels)
            outputs = self.model(images)
            total_losses = []
            for i in range(3):
                loss = self.yolo_losses[i](outputs[i], boxes, labels)
                total_losses.append(loss)
            loss = sum(total_losses)

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            losses.update(loss.item(), images.size(0))
            batch_time.update(time.time() - start)
            start = time.time()

            if i % self.print_freq == 0:
                print('Epoch: [{0}][{1}/{2}]\t'
                      'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Data Time {data_time.val:.3f} ({data_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(epoch, i, len(self.train_loader),
                                                                      batch_time=batch_time,
                                                                      data_time=data_time, loss=losses))
    
    def run(self):
        for epoch in range(self.epochs):
            self.train_epoch(epoch)
            if epoch % self.checkpoint_freq == 0:
                save_state_dict(epoch, self.model, 'yolov3')


if __name__ == '__main__':
    trainer = YoloV3Trainer()
    trainer.run()
